# -*- coding:utf-8 -*-
"""

Author:
    Weichen Shen,wcshen1994@163.com

"""
from __future__ import print_function

import time

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as Data
from torch.optim.lr_scheduler import ReduceLROnPlateau
from sklearn.metrics import *
from torch.utils.data import DataLoader
from tqdm import tqdm

from ..inputs import build_input_features, SparseFeat, DenseFeat, VarLenSparseFeat, get_varlen_pooling_list, \
    create_embedding_matrix
from ..layers import PredictionLayer
from ..layers.utils import slice_arrays

# add
def top1_acc(y_true, y_pred, print_pred=False):
    K=1
    #以下是计算方法
    max_k_preds = y_pred.argsort(axis=1)[:, -K:][:, ::-1] #得到top-k label
    match_array = np.logical_or.reduce(max_k_preds == y_true, axis = 1) #得到匹配结果
    topk_acc_score = match_array.sum() / match_array.shape[0]
    if print_pred:
        pred_labels = sorted(set(y_pred.argmax(axis=-1)))
        print(f"top1_acc\t 分子: {match_array.sum()} \t 分母: {match_array.shape[0]}\t nunique: {len(pred_labels)} \t {pred_labels}")
    return topk_acc_score

def top3_acc(y_true, y_pred):
    K=3
    #以下是计算方法
    max_k_preds = y_pred.argsort(axis=1)[:, -K:][:, ::-1] #得到top-k label
    match_array = np.logical_or.reduce(max_k_preds == y_true, axis = 1) #得到匹配结果
    topk_acc_score = match_array.sum() / match_array.shape[0]
    # print(f"top3_acc\t 分子: {match_array.sum()} \t 分母: {match_array.shape[0]}")
    return topk_acc_score

# origin

class Linear(nn.Module):
    def __init__(self, feature_columns, feature_index, init_std=0.0001, device='cpu', class_num=1):
        super(Linear, self).__init__()
        self.feature_index = feature_index
        self.device = device
        self.sparse_feature_columns = list(
            filter(lambda x: isinstance(x, SparseFeat), feature_columns)) if len(feature_columns) else []
        self.dense_feature_columns = list(
            filter(lambda x: isinstance(x, DenseFeat), feature_columns)) if len(feature_columns) else []

        self.varlen_sparse_feature_columns = list(
            filter(lambda x: isinstance(x, VarLenSparseFeat), feature_columns)) if len(feature_columns) else []

        self.embedding_dict = create_embedding_matrix(feature_columns, init_std, linear=True, sparse=False,
                                                      device=device)

        for tensor in self.embedding_dict.values():
            nn.init.normal_(tensor.weight, mean=0, std=init_std)

        if len(self.dense_feature_columns) > 0:
            print(f"dense_feature_columns is EMPTY!")
            self.weight = nn.Parameter(torch.Tensor(sum(fc.dimension for fc in self.dense_feature_columns), class_num)).to(
                device)
            torch.nn.init.normal_(self.weight, mean=0, std=init_std)
        
        # 多分类
        self.class_num = class_num
        if self.sparse_feature_columns!=[]:
            self.sparse_to_multiclass = nn.Linear(len(self.sparse_feature_columns), self.class_num, bias=False).to(device)
            self.dense_to_multiclass = nn.Linear(len(self.dense_feature_columns), self.class_num, bias=False).to(device)
    

    def forward(self, X):

        sparse_embedding_list = [self.embedding_dict[feat.embedding_name](
            X[:, self.feature_index[feat.name][0]:self.feature_index[feat.name][1]].long()) for
            feat in self.sparse_feature_columns]

        dense_value_list = [X[:, self.feature_index[feat.name][0]:self.feature_index[feat.name][1]] for feat in
                            self.dense_feature_columns]

        varlen_embedding_list = get_varlen_pooling_list(self.embedding_dict, X, self.feature_index,
                                                        self.varlen_sparse_feature_columns, self.device)

        sparse_embedding_list += varlen_embedding_list

        # 增加多分类 linear_sparse_logit linear_dense_logit --> [bsz, num_class]
        if self.class_num>2:
            linear_sparse_logit = self.sparse_to_multiclass(torch.cat(sparse_embedding_list, dim=-1).squeeze())
            linear_dense_logit = torch.cat(dense_value_list, dim=-1).matmul(self.weight)
            linear_logit = linear_sparse_logit + linear_dense_logit
            return linear_logit

        else:
            if len(sparse_embedding_list) > 0 and len(dense_value_list) > 0:
                linear_sparse_logit = torch.sum(
                    torch.cat(sparse_embedding_list, dim=-1), dim=-1, keepdim=False)
                linear_dense_logit = torch.cat(
                    dense_value_list, dim=-1).matmul(self.weight)
                linear_logit = linear_sparse_logit + linear_dense_logit
            elif len(sparse_embedding_list) > 0:
                linear_logit = torch.sum(
                    torch.cat(sparse_embedding_list, dim=-1), dim=-1, keepdim=False)
            elif len(dense_value_list) > 0:
                linear_logit = torch.cat(
                    dense_value_list, dim=-1).matmul(self.weight)
            else:
                linear_logit = torch.zeros([X.shape[0], 1])
            return linear_logit


class BaseModel(nn.Module):
    def __init__(self,
                 linear_feature_columns, dnn_feature_columns, dnn_hidden_units=(128, 128),
                 l2_reg_linear=1e-5,
                 l2_reg_embedding=1e-5, l2_reg_dnn=0, init_std=0.0001, seed=1024, dnn_dropout=0, dnn_activation='relu',
                 task='binary', device='cpu', class_num=1, log_file_path="log.log", model_save_path="model.pt"):

        super(BaseModel, self).__init__()

        self.task = task
        self.dnn_feature_columns = dnn_feature_columns
        self.log_file_path = log_file_path
        self.model_save_path = model_save_path

        self.reg_loss = torch.zeros((1,), device=device)
        self.aux_loss = torch.zeros((1,), device=device)
        self.device = device  # device

        self.feature_index = build_input_features(
            linear_feature_columns + dnn_feature_columns)
        self.dnn_feature_columns = dnn_feature_columns

        self.embedding_dict = create_embedding_matrix(dnn_feature_columns, init_std, sparse=False, device=device)
        #         nn.ModuleDict(
        #             {feat.embedding_name: nn.Embedding(feat.dimension, embedding_size, sparse=True) for feat in
        #              self.dnn_feature_columns}
        #         )

        self.linear_model = Linear(
            linear_feature_columns, self.feature_index, device=device, class_num=class_num)

        self.add_regularization_loss(
            self.embedding_dict.parameters(), l2_reg_embedding)
        self.add_regularization_loss(
            self.linear_model.parameters(), l2_reg_linear)

        self.out = PredictionLayer(task, )
        self.to(device)

    def fit(self, x=None,
            y=None,
            batch_size=None,
            epochs=1,
            verbose=1,
            initial_epoch=0,
            validation_split=0.,
            validation_data=None,
            shuffle=True,
            use_double=False,
            label_weight=None):
        """

        :param x: Numpy array of training data (if the model has a single input), or list of Numpy arrays (if the model has multiple inputs).If input layers in the model are named, you can also pass a
            dictionary mapping input names to Numpy arrays.
        :param y: Numpy array of target (label) data (if the model has a single output), or list of Numpy arrays (if the model has multiple outputs).
        :param batch_size: Integer or `None`. Number of samples per gradient update. If unspecified, `batch_size` will default to 256.
        :param epochs: Integer. Number of epochs to train the model. An epoch is an iteration over the entire `x` and `y` data provided. Note that in conjunction with `initial_epoch`, `epochs` is to be understood as "final epoch". The model is not trained for a number of iterations given by `epochs`, but merely until the epoch of index `epochs` is reached.
        :param verbose: Integer. 0, 1, or 2. Verbosity mode. 0 = silent, 1 = progress bar, 2 = one line per epoch.
        :param initial_epoch: Integer. Epoch at which to start training (useful for resuming a previous training run).
        :param validation_split: Float between 0 and 1. Fraction of the training data to be used as validation data. The model will set apart this fraction of the training data, will not train on it, and will evaluate the loss and any model metrics on this data at the end of each epoch. The validation data is selected from the last samples in the `x` and `y` data provided, before shuffling.
        :param validation_data: tuple `(x_val, y_val)` or tuple `(x_val, y_val, val_sample_weights)` on which to evaluate the loss and any model metrics at the end of each epoch. The model will not be trained on this data. `validation_data` will override `validation_split`.
        :param shuffle: Boolean. Whether to shuffle the order of the batches at the beginning of each epoch.
        :param use_double: Boolean. Whether to use double precision in metric calculation.

        """
        if isinstance(x, dict):
            x = [x[feature] for feature in self.feature_index]
        if validation_data:
            if len(validation_data) == 2:
                val_x, val_y = validation_data
                val_sample_weight = None
            elif len(validation_data) == 3:
                val_x, val_y, val_sample_weight = validation_data  # pylint: disable=unpacking-non-sequence
            else:
                raise ValueError(
                    'When passing a `validation_data` argument, '
                    'it must contain either 2 items (x_val, y_val), '
                    'or 3 items (x_val, y_val, val_sample_weights), '
                    'or alternatively it could be a dataset or a '
                    'dataset or a dataset iterator. '
                    'However we received `validation_data=%s`' % validation_data)
            if isinstance(val_x, dict):
                val_x = [val_x[feature] for feature in self.feature_index]

        elif validation_split and 0. < validation_split < 1.:
            if hasattr(x[0], 'shape'):
                split_at = int(x[0].shape[0] * (1. - validation_split))
            else:
                split_at = int(len(x[0]) * (1. - validation_split))
            x, val_x = (slice_arrays(x, 0, split_at),
                        slice_arrays(x, split_at))
            y, val_y = (slice_arrays(y, 0, split_at),
                        slice_arrays(y, split_at))

        else:
            val_x = []
            val_y = []
        for i in range(len(x)):
            if len(x[i].shape) == 1:
                x[i] = np.expand_dims(x[i], axis=1)

        train_tensor_data = Data.TensorDataset(
            torch.from_numpy(np.concatenate(x, axis=-1)),
            torch.from_numpy(y))
        if batch_size is None:
            batch_size = 256
        train_loader = DataLoader(
            dataset=train_tensor_data, shuffle=shuffle, batch_size=batch_size)

        print(self.device, end="\n")
        model = self.train()
        loss_func = self.loss_func
        optim = self.optim

        sample_num = len(train_tensor_data)
        steps_per_epoch = (sample_num - 1) // batch_size + 1

        info_srt = "Train on {0} samples, validate on {1} samples, {2} steps per epoch".format(
            len(train_tensor_data), len(val_y), steps_per_epoch)
        with open(self.log_file_path, 'a') as f1:
            f1.write(info_srt+"\n")
        print(info_srt)
        scheduler = ReduceLROnPlateau(optim, mode='min', factor=0.5, patience=5, verbose=1)
        for epoch in range(initial_epoch, epochs):
            start_time = time.time()
            loss_epoch = 0
            total_loss_epoch = 0
            # if abs(loss_last - loss_now) < 0.0
            train_result = {}
            loss_epoch_list = []
            
            try:
                with tqdm(enumerate(train_loader), disable=verbose != 1) as t:
                    for index, (x_train, y_train) in t:
                        x = x_train.to(self.device).float()
                        if self.task == 'binary':
                            y = y_train.to(self.device).float()
                        else:
                            y = y_train.to(self.device).long()

                        y_pred = model(x).squeeze()

                        optim.zero_grad()
                        if self.task == 'binary':
                            loss = loss_func(y_pred, y.squeeze(), reduction='sum')
                        elif self.loss_name == "cross_entropy":
                            loss = loss_func(y_pred, y.squeeze(), reduction='mean', ignore_index=0, weight=label_weight)
                        elif self.loss_name == "focal_loss":
                            loss = loss_func(y_pred, y.squeeze())

                        total_loss = loss + self.reg_loss + self.aux_loss

                        loss_epoch += loss.item()
                        total_loss_epoch += total_loss.item()
                        
                        total_loss.backward(retain_graph=True)
                        optim.step()

                        if verbose > 0:
                            for name, metric_fun in self.metrics.items():
                                if name not in train_result:
                                    train_result[name] = []

                                if use_double:
                                    train_result[name].append(metric_fun(
                                        y.cpu().data.numpy(), y_pred.cpu().data.numpy().astype("float64")))
                                else:
                                    train_result[name].append(metric_fun(
                                        y.cpu().data.numpy(), y_pred.cpu().data.numpy()))
                    
                loss_epoch_list.append(loss_epoch)        
                # loss下降 保存模型
                if len(loss_epoch_list) == 1 or loss_epoch_list[-1] < loss_epoch_list[-2]:
                    torch.save(model.state_dict(), self.model_save_path[:-3] + f"_epoch_{epoch}.pt")
                # torch.save(model.state_dict(), self.model_save_path[:-3] + f"_epoch_{epoch}.pt")
                scheduler.step(loss_epoch)
            except KeyboardInterrupt:
                t.close()
                raise
            t.close()

            epoch_time = int(time.time() - start_time)
            if verbose > 0:
                eval_str = "{0}s epoch {1}/{2} - loss: {3: .4f}".format(epoch_time, epoch, epochs, total_loss_epoch / sample_num)

                for name, result in train_result.items():
                    eval_str += " - " + name + ": {0: .4f}".format(np.sum(result) / steps_per_epoch)

                if len(val_x) and len(val_y):
                    eval_result = self.evaluate(val_x, val_y, batch_size)

                    for name, result in eval_result.items():
                        eval_str += " - val_" + name + ": {0: .6f}".format(result)
                print(eval_str)
                with open(self.log_file_path,'a') as f1:
                    f1.write(f'{time.strftime("%Y-%m-%d %H:%M:%S")}\t{eval_str}\n')


    def evaluate(self, x, y, batch_size=256, detail_mode=False):
        """

        :param x: Numpy array of test data (if the model has a single input), or list of Numpy arrays (if the model has multiple inputs).
        :param y: Numpy array of target (label) data (if the model has a single output), or list of Numpy arrays (if the model has multiple outputs).
        :param batch_size:
        :return: Integer or `None`. Number of samples per evaluation step. If unspecified, `batch_size` will default to 256.
        """
        pred_ans = self.predict(x, batch_size)
        eval_result = {}
        for name, metric_fun in self.metrics.items():
            if "acc_top1" in name: 
                eval_result[name] = metric_fun(y, pred_ans, print_pred=True)
            else:
                eval_result[name] = metric_fun(y, pred_ans)

        if self.task == 'binary':
            loss = self.loss_func(torch.from_numpy(pred_ans), torch.from_numpy(y), reduction='sum')
        elif self.loss_name == "cross_entropy":
            loss = self.loss_func(torch.from_numpy(pred_ans), torch.from_numpy(y).squeeze().long(), reduction='mean', ignore_index=0)
        elif self.loss_name == "focal_loss":
            loss = self.loss_func(torch.from_numpy(pred_ans), torch.from_numpy(y).squeeze().long())

        eval_result["loss"] = loss
        if detail_mode: return pred_ans, eval_result
        else: return eval_result

    def predict(self, x, batch_size=256, use_double=False):
        """

        :param x: The input data, as a Numpy array (or list of Numpy arrays if the model has multiple inputs).
        :param batch_size: Integer. If unspecified, it will default to 256.
        :return: Numpy array(s) of predictions.
        """
        model = self.eval()
        if isinstance(x, dict):
            x = [x[feature] for feature in self.feature_index]
        for i in range(len(x)):
            if len(x[i].shape) == 1:
                x[i] = np.expand_dims(x[i], axis=1)

        tensor_data = Data.TensorDataset(
            torch.from_numpy(np.concatenate(x, axis=-1)))
        test_loader = DataLoader(
            dataset=tensor_data, shuffle=False, batch_size=batch_size)

        pred_ans = []
        with torch.no_grad():
            for index, x_test in enumerate(test_loader):
                x = x_test[0].to(self.device).float()
                # y = y_test.to(self.device).float()

                y_pred = model(x).cpu().data.numpy()  # .squeeze()
                pred_ans.append(y_pred)

        if use_double:
            return np.concatenate(pred_ans).astype("float64")
        else:
            return np.concatenate(pred_ans)

    def input_from_feature_columns(self, X, feature_columns, embedding_dict, support_dense=True):

        sparse_feature_columns = list(
            filter(lambda x: isinstance(x, SparseFeat), feature_columns)) if len(feature_columns) else []
        dense_feature_columns = list(
            filter(lambda x: isinstance(x, DenseFeat), feature_columns)) if len(feature_columns) else []

        varlen_sparse_feature_columns = list(
            filter(lambda x: isinstance(x, VarLenSparseFeat), feature_columns)) if feature_columns else []

        if not support_dense and len(dense_feature_columns) > 0:
            raise ValueError(
                "DenseFeat is not supported in dnn_feature_columns")

        sparse_embedding_list = [embedding_dict[feat.embedding_name](
            X[:, self.feature_index[feat.name][0]:self.feature_index[feat.name][1]].long()) \
                for feat in sparse_feature_columns]

        varlen_sparse_embedding_list = get_varlen_pooling_list(self.embedding_dict, X, self.feature_index,
                                                               varlen_sparse_feature_columns, self.device)

        dense_value_list = [X[:, self.feature_index[feat.name][0]:self.feature_index[feat.name][1]] for feat in
                            dense_feature_columns]

        return sparse_embedding_list + varlen_sparse_embedding_list, dense_value_list

    def compute_input_dim(self, feature_columns, include_sparse=True, include_dense=True, feature_group=False):
        sparse_feature_columns = list(
            filter(lambda x: isinstance(x, (SparseFeat, VarLenSparseFeat)), feature_columns)) if len(
            feature_columns) else []
        dense_feature_columns = list(
            filter(lambda x: isinstance(x, DenseFeat), feature_columns)) if len(feature_columns) else []

        dense_input_dim = sum(
            map(lambda x: x.dimension, dense_feature_columns))
        if feature_group:
            sparse_input_dim = len(sparse_feature_columns)
        else:
            sparse_input_dim = sum(feat.embedding_dim for feat in sparse_feature_columns)
        input_dim = 0
        if include_sparse:
            input_dim += sparse_input_dim
        if include_dense:
            input_dim += dense_input_dim
        return input_dim

    def add_regularization_loss(self, weight_list, weight_decay, p=2):
        reg_loss = torch.zeros((1,), device=self.device)
        for w in weight_list:
            if isinstance(w, tuple):
                l2_reg = torch.norm(w[1], p=p, )
            else:
                l2_reg = torch.norm(w, p=p, )
            reg_loss = reg_loss + l2_reg
        reg_loss = weight_decay * reg_loss
        self.reg_loss = self.reg_loss + reg_loss

    def add_auxiliary_loss(self, aux_loss, alpha):
        self.aux_loss = aux_loss * alpha

    def compile(self, optimizer,
                loss=None,
                metrics=None,
                alpha=None
                ):
        """
        :param optimizer: String (name of optimizer) or optimizer instance. See [optimizers](https://pytorch.org/docs/stable/optim.html).
        :param loss: String (name of objective function) or objective function. See [losses](https://pytorch.org/docs/stable/nn.functional.html#loss-functions).
        :param metrics: List of metrics to be evaluated by the model during training and testing. Typically you will use `metrics=['accuracy']`.
        """

        self.optim = self._get_optim(optimizer)
        self.loss_func = self._get_loss_func(loss,alpha=alpha)
        self.metrics = self._get_metrics(metrics)

    def _get_optim(self, optimizer):
        if isinstance(optimizer, str):
            if optimizer == "sgd":
                optim = torch.optim.SGD(self.parameters(), lr=0.01)
            elif optimizer == "adam":
                optim = torch.optim.Adam(self.parameters())  # 0.001
            elif optimizer == "adagrad":
                optim = torch.optim.Adagrad(self.parameters())  # 0.01
            elif optimizer == "rmsprop":
                optim = torch.optim.RMSprop(self.parameters())
            else:
                raise NotImplementedError
        else:
            optim = optimizer
        return optim

    def _get_loss_func(self, loss, alpha=None):
        if isinstance(loss, str):
            if loss == "binary_crossentropy":
                loss_func = F.binary_cross_entropy
            elif loss == "mse":
                loss_func = F.mse_loss
            elif loss == "mae":
                loss_func = F.l1_loss
            elif loss == "cross_entropy":
                loss_func = F.cross_entropy
            elif loss == "focal_loss":
                from .Focal_Loss import FocalLoss
                loss_func = FocalLoss(class_num=122, alpha=alpha, gamma=2, use_alpha=True)
            else:
                raise NotImplementedError
        else:
            loss_func = loss
        self.loss_name = loss
        return loss_func

    def _log_loss(self, y_true, y_pred, eps=1e-7, normalize=True, sample_weight=None, labels=None):
        # change eps to improve calculation accuracy
        return log_loss(y_true,
                        y_pred,
                        eps,
                        normalize,
                        sample_weight,
                        labels)

    def _get_metrics(self, metrics, set_eps=False):
        metrics_ = {}
        if metrics:
            for metric in metrics:
                if metric == "binary_crossentropy" or metric == "logloss":
                    if set_eps:
                        metrics_[metric] = self._log_loss
                    else:
                        metrics_[metric] = log_loss
                if metric == "auc":
                    metrics_[metric] = roc_auc_score
                if metric == "mse":
                    metrics_[metric] = mean_squared_error
                if metric == "accuracy" or metric == "acc" or metric == "acc_top1":
                    if self.task == 'binary':
                        metrics_[metric] = lambda y_true, y_pred: accuracy_score(y_true, np.where(y_pred > 0.5, 1, 0))
                    else:
                        # metrics_[metric] = lambda y_true, y_pred: accuracy_score(y_true, y_pred.argmax(axis=-1))
                        metrics_[metric] = top1_acc
                if metric == "acc_top3":
                    metrics_[metric] = top3_acc
        return metrics_

    @property
    def embedding_size(self, ):
        feature_columns = self.dnn_feature_columns
        sparse_feature_columns = list(
            filter(lambda x: isinstance(x, (SparseFeat, VarLenSparseFeat)), feature_columns)) if len(
            feature_columns) else []
        embedding_size_set = set([feat.embedding_dim for feat in sparse_feature_columns])
        if len(embedding_size_set) > 1:
            raise ValueError("embedding_dim of SparseFeat and VarlenSparseFeat must be same in this model!")
        return list(embedding_size_set)[0]

